'''
Created on Apr 28, 2011

@author: oabalbin
'''
import sys
import numpy as np
from optparse import OptionParser
from exome.vcftools.vcf import *

def write_vcf_line(snp, infotags, sampletags):
    '''
    '''
    site=snp['COORD'].split('@')
    CHROM, POS=str(site[0]),str(site[1])
    
    qinfo,asinfo=deque(),deque()
    for tag in infotags:
        if tag not in snp.keys():
            qinfo.append(tag+'=$$')
            continue
        if tag=='CI95' or tag=='SF':
            b=tag+'='
            for j,a in enumerate(snp[tag]):
                if j==0:
                    b+=str(a)
                else:
                    b+='&'+str(a)
            qinfo.append(b)
        else:
            qinfo.append(tag+'='+str(snp[tag]))
    
    qinfo = [",".join(qinfo).replace(',',';')]
    # snps samples
    samples = snp['_SAMPLES']
    gt={'ref':'0','alt':'1'}
    #print samples
    for sp in samples:
        sinfo=deque()
        #print sp
        for tag in sampletags:
            if tag=="GT":
                gtref=gt[sp['GT1']]
                gtalt=gt[sp['GT2']]
                sinfo.append(gtref+'/'+gtalt)
                continue
            elif tag=="AD":
                ADref=str(sp['ADref'])
                ADalt=str(sp['ADalt'])
                sinfo.append(ADref+'&'+ADalt)
                continue
            
            elif tag=="GL":
                GLrr = str(sp['GLrr'])
                GLra = str(sp['GLra'])
                GLaa = str(sp['GLaa'])
                sinfo.append(GLrr+'&'+GLra+'&'+GLaa)
                continue
            
            sinfo.append(str(sp[tag]))
        asinfo.append((",".join(sinfo).replace(',',':')))
    
    sampletags = [",".join(sampletags).replace(',',':')]
                    
    gralinfo = [CHROM, 
            POS,
            snp['ID'], 
            snp['REF'], 
            snp['ALT'], 
            str(snp['QUAL']), 
            snp['FILTER']]
    
    line = gralinfo + qinfo + sampletags + list(asinfo)
    return ",".join(line).replace(',','\t').replace('&',',')
    

def snps_het(vcfObject, vcfOutfile=False, xsamp=1):
    '''
    Extract heterozygous snvs from a vcf file. 
    Write a vcf file with the extracted snps
    '''
    hetsnps=0
    heterozygocity_low_th=0.45
    heterozygocity_high_th=0.65
    
    for snp in vcfObject:
        samples = snp['_SAMPLES']
        het=np.zeros(len(samples))
        
        # Filter according to PASS filter and AF
        if (snp['FILTER']=="PASS" or snp['FILTER']==".") and \
            (snp['AF'] >= heterozygocity_low_th and snp['AF'] <= heterozygocity_high_th): 
            i=0
            if i==0 and (len(samples) <= xsamp):
                print >> sys.stderr,"Total Number of samples is smaller than intersection samples %d"%xsamp
                print >> sys.stderr,"Total Number of samples %d"%len(samples)
                print snp
            for sp in samples:
                if (sp['GT1']=='ref' and sp['GT2']=='alt'):
                    het[i]=1        
                i+=1
            print len(samples),sum(het), (len(samples) - xsamp)
            if sum(het)==(len(samples) - xsamp):
                line = write_vcf_line(snp, vcfObject.infotags['_INFOTAGS'], vcfObject.infotags['_SAMPLETAGS'])
                hetsnps+=1
                #print line
                if vcfOutfile:
                    vcfOutfile.write(line+'\n')
                
    print >> sys.stderr,"Total Heterozygous SNPs %d"%hetsnps


def snps_bySample(vcfObject, vcfOutfile=False, thissample=None):
    '''
    Extract heterozygous snvs from a vcf file. 
    Write a vcf file with the extracted snps
    '''
    hetsnps=0
    
    for snp in vcfObject:
        samples = snp['_SAMPLES']
        print samples
        # Filter according to PASS filter and AF
        if (snp['FILTER']=="PASS" or snp['FILTER']=="."):
            print samples
            for sp in samples:
                #print sp['NAME'],thissample, sp['GT1'], sp['GT2']
                if sp['NAME']== thissample and (sp['GT1']!='ref' and sp['GT2']!='ref'):
                    print sp
                    line = write_vcf_line(snp, vcfObject.infotags['_INFOTAGS'], vcfObject.infotags['_SAMPLETAGS'])
                    hetsnps+=1
                    #print line
                    if vcfOutfile:
                        vcfOutfile.write(line+'\n')
                
    print >> sys.stderr,"Total Heterozygous SNPs %d"%hetsnps




def extract_heterozygous_snps(input_file_name, vcf_outfile_name,xsamp):
    '''
    '''
    vcOutfile=open(vcf_outfile_name,'w')
    thisvcf=VCF()
    thisvcf.readFromFile(input_file_name)
    for line in thisvcf.header:
        vcOutfile.write(line+'\n')
    snps_het(thisvcf,vcOutfile,xsamp)
    vcOutfile.close()


def extract_snps_sample(input_file_name, vcf_outfile_name, patient_name):
    vcOutfile=open(vcf_outfile_name,'w')
    thisvcf=VCF()
    thisvcf.readFromFile(input_file_name)
    for line in thisvcf.header:
        vcOutfile.write(line+'\n')
    snps_bySample(thisvcf,vcOutfile, patient_name)
    vcOutfile.close()



if __name__ == '__main__':
        
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)
    optionparser.add_option("-f", "--vcf_file", dest="vcf_file",
                            help="vcf_file")
    optionparser.add_option("-o", "--out_vcf_file", dest="out_vcf_file",
                            help="out_vcf_file")
    optionparser.add_option("-x", "--samples", action='append', dest="samples",
                            help="out_vcf_file")
    
    
    (options, args) = optionparser.parse_args()
    samples = options.samples
    print samples
    for sample in samples:
        extract_snps_sample(options.vcf_file, options.out_vcf_file+'_'+sample,sample)
        